from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt


plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


m_A = 4866
m_B = 2433
f = 6250
delta_m = 1335.535
omiga = 1.4005
k_3 = 656.3616  # 兴波
k_5 = 80000  # 弹簧劲度系数
k_6 = 10000  # 阻尼系数
ro = 1025
g = 9.8
s = np.pi


# 定义一个方程组（微分方程组）
def pfun(y, t):
    y0, y1, y2, y3 = y
    return np.array([
        1 / (m_A + delta_m) * (-k_3 * y0 - ro * g * s * y1 + f * np.cos(omiga * t) + k_5 * (y3 - y1) + k_6 * (y2 - y0)),
        y0,
        1 / m_B * (-k_5 * (y3 - y1) - k_6 * (y2 - y0)),
        y2
    ])


t = np.arange(0, 28.56122813280971, 0.2)  # 创建自变量序列
soli = odeint(pfun, [0.0, 0.0, 0.0, 0.0], t)  # 求数值解
plt.rc('font', size=16)
plt.rc('font', family='SimHei')
plt.plot(t, soli[:, 1], 'b*', label="浮子位移")
plt.plot(t, soli[:, 3], 'r*', label="振子位移")
plt.legend()
plt.show()
print(soli[50, 1])
